import numpy as np
from modules.feature_extracter_without_delta_layer import featureExtracter


checkpoint = torch.load("./amodel.pth.tar")
amodel = featureExtracter(channels=1)
amodel.load_state_dict(checkpoint['state_dict']) 
amodel.cuda()
amodel.eval()

example = torch.rand(1, 1, 32, 900)
example = example.cuda()
torch.onnx.export(amodel,
                 (example),
                 'overlapTransformer.onnx',
                 input_names = ['input'],
                 output_names = ['output'],
                 opset_version=11,
                 verbose = True)
